import torch
import torch.nn as nn
import torchvision.models as models

class FruitClassifier(nn.Module):
    def __init__(self, num_classes=30, model_name='resnet50', dropout=0.2):
        """
        初始化分类器
        Args:
            num_classes: 类别数量
            model_name: 使用的预训练模型名称
                - 'resnet50': ResNet50
                - 'mobilenet_v3_small': MobileNetV3 Small
                - 'efficientnet_b0': EfficientNet-B0
            dropout: dropout率
        """
        super(FruitClassifier, self).__init__()
        
        # 加载预训练模型
        if model_name == 'resnet50':
            self.base_model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
            num_features = self.base_model.fc.in_features
            self.base_model.fc = nn.Identity()
        elif model_name == 'mobilenet_v3_small':
            self.base_model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)
            num_features = self.base_model.classifier[-1].in_features
            self.base_model.classifier[-1] = nn.Identity()
        elif model_name == 'efficientnet_b0':
            self.base_model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)
            num_features = self.base_model.classifier[-1].in_features
            self.base_model.classifier[-1] = nn.Identity()
        
        # 创建分类头
        self.classifier = nn.Sequential(
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(512, num_classes)
        )
        
        # 默认冻结特征提取器
        for param in self.base_model.parameters():
            param.requires_grad = False
            
    def forward(self, x):
        # 特征提取
        features = self.base_model(x)
        # 分类
        output = self.classifier(features)
        return output

    def unfreeze_features(self, unfreeze_last_n_layers=2):
        """解冻最后几层特征提取层进行微调"""
        if isinstance(self.base_model, models.ResNet):
            # ResNet的层
            layers_to_unfreeze = [
                self.base_model.layer4,
                self.base_model.layer3,
            ][:unfreeze_last_n_layers]
        elif isinstance(self.base_model, models.MobileNetV3):
            # MobileNetV3的层
            layers_to_unfreeze = [
                self.base_model.features[-1],
                self.base_model.features[-2],
            ][:unfreeze_last_n_layers]
        elif isinstance(self.base_model, models.EfficientNet):
            # EfficientNet的层
            layers_to_unfreeze = [
                self.base_model.features[-1],
                self.base_model.features[-2],
            ][:unfreeze_last_n_layers]
            
        # 解冻选定的层
        for layer in layers_to_unfreeze:
            for param in layer.parameters():
                param.requires_grad = True